import numcodecs
import zarr
from pathlib import Path
import numpy as np
import yaml
from pynput import keyboard
import time
from ezrobot.logger.data_logger import TrajectoryDataLogger

from ezrobot.utils import GraphRunner, HandEyeConfig
from ezrobot.sensor import ViveTracker, SensorManager, AngleEncoder, RealSense, TactileSensor
from ezrobot.utils.general import accurate_sleep
from ezgl import tb, Matrix4x4
from enum import Enum



BASE_DIR = Path(__file__).resolve().parent
CONFIG_DIR = BASE_DIR.parent / "configs"

class RunnerState(Enum):
    Idle = 0
    Running = 1
    Stopped = 2
    StartRecording = 3
    StopRecording = 4

 
  
   
    
     
      
       
        
         
          
           
            
             
              
               
                
                 
                  
                   
                    
                     
                      
                       

class UMIRunner:
    def __init__(self, frequency=50):
        # init flags
        self.frequency = frequency
        self.frequency_s = 1/self.frequency
        self.state = RunnerState.Idle

        # init logger
        self.logger = TrajectoryDataLogger(path=BASE_DIR / "bread_data.zarr", start_id=127, overwrite=False)

        # init keyboard listener
        listener = keyboard.Listener(on_press=self.on_key_press)
        listener.start()    
        
        # init utils
        self.ee2gripper = np.array([[1,0,0,0],
               [0,1,0,0],
               [0,0,1,0.11772],
               [0,0,0,1]])
        # self.compressor = numcodecs.Blosc(cname='zstd', clevel=5, shuffle=numcodecs.Blosc.SHUFFLE)
        self.handeye_config = HandEyeConfig.load(CONFIG_DIR / "handeye_config.yaml")
        with open(CONFIG_DIR/"config.yaml", 'r') as f:
                self.project_config = yaml.safe_load(f)

        # init sensors 
        self.sensor_encoder = AngleEncoder(self.project_config['UMI']["encoder_id"], self.project_config['UMI']["encoder_index"])
        self.sensor_vive = ViveTracker(1)
        self.sensor_tactile_UMI_l = TactileSensor(self.project_config['UMI']['tactile_ser_num']['left'])
        self.sensor_rs = RealSense(index=self.project_config['robot_settings']['realsense_ser_num'])
        self.sensor_manager = SensorManager(sensor=[self.sensor_encoder, self.sensor_vive, self.sensor_tactile_UMI_l, self.sensor_rs ], 
                                            frequency=30, 
                                            buffer_size=10, 
                                            streaming=True)
    def on_key_press(self, key):
        if key == keyboard.Key.space:
            self.state = RunnerState.StopRecording
        elif key == keyboard.Key.enter:
            self.state = RunnerState.StartRecording
        elif key == keyboard.Key.esc:
            self.state = RunnerState.Stopped


    def get_target_pose(self, vive_pose):
        """
        获取目标位姿, vive2hand
        """
        return self.handeye_config.base2map * vive_pose * self.handeye_config.hand2eye.inverse() * Matrix4x4(self.ee2gripper)


    def run(self):
        print("press 'Enter' to start recording")
        while True:    
            t_start = time.monotonic()
            if self.state == RunnerState.Stopped:
                break

            elif self.state == RunnerState.Idle:
                pass

            elif self.state == RunnerState.Running:
                data = self.sensor_manager.get_data()
                target_pose = self.get_target_pose(data['ViveTracker_1'])
                self.logger.record_frame(
                    data[f"RealSense_{self.project_config['robot_settings']['realsense_ser_num']}"]['color'][:250,:].copy(),
                    data[f"RealSense_{self.project_config['robot_settings']['realsense_ser_num']}"]['depth'][:250,:].copy(),
                    data['AngleEncoder_1'],
                    target_pose.toTransRotVector(),
                    data[f"TactileSensor_{self.project_config['UMI']['tactile_ser_num']['left']}"]['src'],
                    data[f"TactileSensor_{self.project_config['UMI']['tactile_ser_num']['left']}"]['force']
                )

            elif self.state == RunnerState.StartRecording:
                print("################ Start Recording, press 'Space' to stop ################")
                self.logger.start_trajectory()
                self.state = RunnerState.Running

            elif self.state == RunnerState.StopRecording:
                print("################ Stop Recording, press 'Enter' to start ################")
                self.logger.end_trajectory()
                self.state = RunnerState.Idle

            # sleep accurately
            elapsed_time = time.monotonic() - t_start
            if self.frequency_s > elapsed_time:
                accurate_sleep(self.frequency_s - elapsed_time)    
            
if __name__ == "__main__":
    project_runner = UMIRunner()
    project_runner.run()